#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import torch
import torch.nn as nn
import math
import torch.nn.init as init
import re,sys,os
from . import resnet
#import resnet

class PFLDResNet(nn.Module):
    def __init__(self):
        super(PFLDResNet, self).__init__()
        layers=[2,2,2,5]
        strides=[2,2,2,1]
        channels=[64,128,256,256]
        '''
        self.backbone = resnet.ResNet(layers=layers, strides=strides, channels=channels, block=resnet.Bottleneck).model
        self.ldmk_head = nn.Conv2d(2048, 1024, kernel_size=3, padding=1, stride=1)
        self.flatten = nn.Flatten()
        self.ldmk_fc = nn.Linear(4*4*1024, 2*98)
        self.pose_head = nn.Conv2d(2048, 256, kernel_size=3, padding=1, stride=1)
        self.pose_fc = nn.Linear(4*4*256, 3)
        '''
        sys.path.append('/home/jianjun/workspace/generic-detector')
        from vision.ssd import myresnet

        resnet = myresnet.ResNet(layers=layers, strides=strides, channels=channels, block=myresnet.ResidualBlock)
        self.backbone = resnet.model
        self.flatten = nn.Flatten()
        self.ldmk_head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(7*7*256, 2*98),
        )
        self.pose_head = nn.Sequential(
            resnet.make_layer(resnet.block, 256, 128, 5),
            nn.Flatten(),
            nn.Linear(7*7*128, 3)
        )
 
    def forward(self, x):  # x: 3, 112, 112
        for k, layer in self.backbone.named_children():
            if k == 'layer4':
                branch1 = x
            x = layer(x)
                
        ldmk = x
        ldmk = self.ldmk_head(ldmk)

        pose = branch1
        pose = self.pose_head(pose)

        return pose, ldmk


if __name__ == '__main__':
    input = torch.randn(1, 3, 112, 112)
    plfd_backbone = PFLDResNet()
    from torchstat import stat
    stat(plfd_backbone, (3,112,112))
    angle, landmarks = plfd_backbone(input)
    print(plfd_backbone)
    print("angle.shape:{0:}, landmarks.shape: {1:}".format(
        angle.shape, landmarks.shape))
